module mat_mul_oper
    # (parameter ADDR_WIDTH = 4,
       parameter DATA_WIDTH = 32,
       parameter DEPTH = 16
    )
    (
        input                       clk,
        input                       rst_n,
        input                       kick_start,
        output reg                  ready,
        // 用户接口 - 简化操作
        input      [ADDR_WIDTH-1:0] user_addr,
        input      [DATA_WIDTH-1:0] user_data_in,
        output reg [DATA_WIDTH-1:0] user_data_out,
        input                       kick_write,
        input                       kick_read,
        output reg [DATA_WIDTH-1:0] matrix_a_output [0:3],
        output reg [DATA_WIDTH-1:0] matrix_b_output [0:3],
        output reg [DATA_WIDTH-1:0] matrix_c_output [0:3],
        output reg [DATA_WIDTH-1:0] matrix_p_output
    );

    // 接收外部的kick信号
    reg kick_start_triggered;
    reg kick_write_triggered;
    reg kick_read_triggered;
    // 状态机
    reg [7:0] state_machine;
    localparam MOD_BOOT                   = 4'h0;
    localparam READ_PRE                   = 4'h1;
    localparam READ_EXEC                  = 4'h2;
    localparam READ_END                   = 4'h3;
    localparam WRITE_PRE                  = 4'h4;
    localparam WRITE_EXEC_DONE            = 4'h5;
    localparam MAC_START                  = 4'h6;
    localparam MAC_CACHE_LOAD_MAT_A_START = 4'h7;
    localparam MAC_CACHE_LOAD_MAT_B_START = 4'h8;
    localparam MAC_CALC_MAT_C             = 4'h9;
    localparam MAC_READ_AB_END            = 4'hA;
    localparam MOD_FUCKING_SLEEP          = 4'hB;
    localparam MAC_CACHE_LOAD_MAT_A_END   = 4'hC;
    localparam MAC_CACHE_LOAD_MAT_B_END   = 4'hD;
    localparam MAC_OUTPUT_MAT_C           = 4'hE;
    localparam MOD_POOLING_START          = 4'hF;
    localparam MOD_POOLING_CALC           = 4'h10;
    localparam MOD_POOLING_END            = 4'h11;
    // 连接到原始RAM模块的信号
    reg  [ADDR_WIDTH-1:0] ram_addr;
    wire [DATA_WIDTH-1:0] ram_data;
    reg                   ram_cs;
    reg                   ram_we;
    reg                   ram_oe;
    // 控制对四个矩阵的操作
    reg [1:0] matrix_access_flag;
    localparam MATRIX_A_READ  = 2'b00;
    localparam MATRIX_B_READ  = 2'b01;
    localparam MATRIX_C_WRITE = 2'b10;
    localparam MATRIX_NONE    = 2'b11;

    // 三态数据线处理
    reg [DATA_WIDTH-1:0] ram_write_data_reg;
    reg                  ram_write_en; 
    // 开启ram write信号后，ram模块的data信号将作为input信号，接收ram_write_data_reg输入的信号，进行写操作,而当关闭ram write信号后，ram模块的data将作为outpupt信号，进行读操作
    assign ram_data = ram_write_en ? ram_write_data_reg : {DATA_WIDTH{1'bz}};

    // 用于MAC计算的相关处理
    reg [3:0] vector_id;
    reg [1:0] mac_cnt;
    reg [3:0] cur_matrix_base_addr;
    localparam MATRIX_A_BASE_ADDR = 4'h0;
    localparam MATRIX_B_BASE_ADDR = 4'h4;
    localparam MATRIX_C_BASE_ADDR = 4'h8;
    localparam MATRIX_P_BASE_ADDR = 4'hc;

    reg [DATA_WIDTH-1:0]  matrix_a_cache;
    reg [DATA_WIDTH-1:0]  matrix_b_cache;

    reg [DATA_WIDTH-1:0]  matrix_a [0:3];
    reg [DATA_WIDTH-1:0]  matrix_b [0:3];
    reg [DATA_WIDTH-1:0]  matrix_c [0:3];
    reg [DATA_WIDTH-1:0]  matrix_p;
    // 输出matrix_a matrix_b matrix_c信号

    /* Mac operation function set */
    function [7:0] get_matrix_element;
        input [DATA_WIDTH-1:0] vector;
        input [1:0]            vector_id;
        begin
            case (vector_id)
                2'b00: get_matrix_element = vector[31:24];
                2'b01: get_matrix_element = vector[23:16];
                2'b10: get_matrix_element = vector[15: 8];
                2'b11: get_matrix_element = vector[7 : 0];
                default: 
                    get_matrix_element = 0;         
            endcase
        end
    endfunction

    function [15:0] product;
        input [7:0] element_vector_a;
        input [7:0] element_vector_b;
        begin
            product = element_vector_a * element_vector_b;
        end
    endfunction

    function [7:0] single_vector_mac;
        input [31:0] vector_a;
        input [31:0] vector_b;
        reg   [15:0] single_mac_ret;
        reg   [ 7:0] element_a;
        reg   [ 7:0] element_b;
        integer i;
        begin
            single_mac_ret = 0;
            for (i = 0; i < 4; i = i + 1) begin
                element_a = get_matrix_element(vector_a, i[1:0]);
                element_b = get_matrix_element(vector_b, i[1:0]);
                single_mac_ret = single_mac_ret + product(element_a, element_b);
            end     
            single_vector_mac = single_mac_ret;        
        end
    endfunction
    /* Mac operation function set */

    /* Mac calc combinational circuit */
    function void matrix_mac;
        input  [31:0] matrix_a_input     [0:3];
        input  [31:0] matrix_b_input     [0:3];
        output [31:0] matrix_c_output    [0:3]; 
        reg    [ 7:0] element_vector_c0;
        reg    [ 7:0] element_vector_c1;
        reg    [ 7:0] element_vector_c2;
        reg    [ 7:0] element_vector_c3;
        integer i;
        begin
            for (i = 0; i < 4; i = i + 1) begin
                element_vector_c0  = single_vector_mac(matrix_a_input[i], matrix_b_input[0]);
                element_vector_c1  = single_vector_mac(matrix_a_input[i], matrix_b_input[1]);
                element_vector_c2  = single_vector_mac(matrix_a_input[i], matrix_b_input[2]);
                element_vector_c3  = single_vector_mac(matrix_a_input[i], matrix_b_input[3]);
                matrix_c_output[i] = {element_vector_c0, element_vector_c1, element_vector_c2, element_vector_c3};
                $display("matrix_c_output[%0d] = 0x%h", i, matrix_c_output[i]);
            end 
        end
    endfunction

    /* Pooling operation function set */
    function void matrix_pooling;
        input  [31:0] matrix_c_i [0:3];
        output [31:0] matrix_p_o;
        reg    [ 7:0] element_matrix_p00;
        reg    [ 7:0] element_matrix_p01;
        reg    [ 7:0] element_matrix_p10;
        reg    [ 7:0] element_matrix_p11;
        reg    [ 7:0] temp_element_00;
        reg    [ 7:0] temp_element_01;
        reg    [ 7:0] temp_element_10;
        reg    [ 7:0] temp_element_11;
        integer i;
        begin
            /* p00 */
            temp_element_00    = get_matrix_element(matrix_c_i[0], 2'b00);
            temp_element_01    = get_matrix_element(matrix_c_i[0], 2'b01);
            temp_element_10    = get_matrix_element(matrix_c_i[1], 2'b00);
            temp_element_11    = get_matrix_element(matrix_c_i[1], 2'b01);
            element_matrix_p00 = (temp_element_00 + temp_element_01 + temp_element_10 + temp_element_11) / 4;
            /* p01 */
            temp_element_00    = get_matrix_element(matrix_c_i[0], 2'b10);
            temp_element_01    = get_matrix_element(matrix_c_i[0], 2'b11);
            temp_element_10    = get_matrix_element(matrix_c_i[1], 2'b10);
            temp_element_11    = get_matrix_element(matrix_c_i[1], 2'b11);
            element_matrix_p01 = (temp_element_00 + temp_element_01 + temp_element_10 + temp_element_11) / 4;
            /* p10 */
            temp_element_00    = get_matrix_element(matrix_c_i[2], 2'b00);
            temp_element_01    = get_matrix_element(matrix_c_i[2], 2'b01);
            temp_element_10    = get_matrix_element(matrix_c_i[3], 2'b00);
            temp_element_11    = get_matrix_element(matrix_c_i[3], 2'b01);
            element_matrix_p10 = (temp_element_00 + temp_element_01 + temp_element_10 + temp_element_11) / 4;
            /* p11 */
            temp_element_00    = get_matrix_element(matrix_c_i[2], 2'b10);
            temp_element_01    = get_matrix_element(matrix_c_i[2], 2'b11);
            temp_element_10    = get_matrix_element(matrix_c_i[3], 2'b10);
            temp_element_11    = get_matrix_element(matrix_c_i[3], 2'b11);
            element_matrix_p11 = (temp_element_00 + temp_element_01 + temp_element_10 + temp_element_11) / 4;
            matrix_p_o         = {element_matrix_p00, element_matrix_p01, element_matrix_p10, element_matrix_p11};
            $display("matrix_p_output[0] = 0x%h", matrix_p_o[31:16]);
            $display("matrix_p_output[1] = 0x%h", matrix_p_o[15: 0]);
            $display("matrix_p_output    = 0x%h", matrix_p_o);
        end
    endfunction
    /* Pooling operation function set */

    always @(posedge clk or negedge rst_n) begin
        if (!rst_n) begin
            state_machine      <= MOD_FUCKING_SLEEP;
            ready              <= 1'b1;
            ram_addr           <= 4'h0;
            // 只要cs没开，ram模块就不工作
            ram_cs             <= 1'b0;
            ram_we             <= 1'b0;
            ram_oe             <= 1'b0;
            ram_write_en       <= 1'b0;
            ram_write_data_reg <= 32'h00000000;
        end else begin
            case (state_machine)
                MOD_FUCKING_SLEEP: begin
                    //$display("kimi1=====MAC_SLEEP====i am fucking sleep...");
                    if (kick_write && ready) begin
                        kick_write_triggered <= 1'b1;
                        state_machine        <= MOD_BOOT;
                        ready                <= 1'b0;
                        $display("\n[kick_write]************* write mem boot...");
                    end else if (kick_read && ready) begin
                        kick_read_triggered  <= 1'b1;
                        state_machine        <= MOD_BOOT;
                        ready                <= 1'b0;
                        $display("\n[kick_read]************* read mem boot...");
                    end else if (kick_start && ready) begin
                        kick_start_triggered <= 1'b1;
                        state_machine        <= MOD_BOOT;
                        ready                <= 1'b0;
                        $display("\n[kick_start]************* mac boot...");
                    end else begin
                        state_machine        <= MOD_FUCKING_SLEEP;
                        ready                <= 1'b1;
                        //$display("[MAC_SLEEP]=========i am fucking sleep...kick_start %d, ready %d", kick_start, ready);
                    end
                end
                MOD_BOOT: begin
                    //$display("kimi=====MOD_BOOT");
                    if (kick_write_triggered) begin
                        state_machine <= WRITE_PRE;
                    end else if (kick_read_triggered) begin
                        state_machine <= READ_PRE;
                    end else if (kick_start_triggered) begin
                        //$display("kimi=====kick_start_triggered");
                        state_machine <= MAC_START;
                        vector_id     <= 4'h0;
                    end else begin
                        state_machine <= MOD_FUCKING_SLEEP;
                        ready         <= 1'b1;
                    end
                end
                /****************************** Write Operation ************************************/
                WRITE_PRE: begin
                        //$display("kimi=====WRITE_PRE");
                        ram_addr             <= user_addr;
                        ram_write_en         <= 1'b1;
                        ram_write_data_reg   <= user_data_in;
                        ram_cs               <= 1'b1;
                        ram_we               <= 1'b1;
                        ram_oe               <= 1'b0;
                        state_machine        <= WRITE_EXEC_DONE;
                        kick_write_triggered <= 1'b0;
                        //$display("[WRITE_PRE] input: addr: 0x%h data: 0x%h", user_addr, user_data_in);
                end
                WRITE_EXEC_DONE: begin
                        //$display("kimi=====WRITE_EXEC_DONE");
                        ram_addr             <= 4'h0;
                        ram_write_en         <= 1'b0;
                        ram_write_data_reg   <= 32'h00000000;;
                        ram_cs               <= 1'b0;
                        ram_we               <= 1'b0;
                        ram_oe               <= 1'b0;
                        state_machine        <= MOD_FUCKING_SLEEP;
                        ready                <= 1'b1;
                        kick_write_triggered <= 1'b0;
                        //$display("[WRITE_EXEC_DONE] exec mem write: addr 0x%h data 0x%h ram_write_en %d ram_cs %d ram_we %d ram_oe %d  ready %d",
                            //ram_addr, ram_write_data_reg, ram_write_en, ram_cs, ram_we, ram_oe, ready);
                end
                /****************************** Read Operation ************************************/
                READ_PRE: begin
                    //$display("\n*************kick_start: %d", kick_start);
                    state_machine <= READ_EXEC;
                    if (kick_start_triggered) begin
                        ram_addr  <= cur_matrix_base_addr + vector_id;
                    end else begin
                        ram_addr  <= user_addr;
                    end
                    // cs开启，ram模块开始工作，oe=1,we=0 通知ram进行读操作
                    ram_cs        <= 1'b1;
                    ram_we        <= 1'b0;
                    ram_oe        <= 1'b1;
                    //$display("\n[READ_PRE]************* addr: 0x%h cs: %d we: %d oe: %d kick_start_triggered %d kick_start %d cur_matrix_base_addr %d vector_id %d",
                        //ram_addr, ram_cs, ram_we, ram_oe, kick_start_triggered, kick_start, cur_matrix_base_addr, vector_id);
                end
                READ_EXEC: begin
                    //$display("kimi=====READ_EXEC");
                    // 当前这个posedge ram那边正在读出数据，等待一个周期，让RAM输出稳定，下一个posedge拿数据
                    state_machine <= READ_END;
                    //$display("\n[READ_EXEC]************* addr: 0x%h cs: %d we: %d oe: %d kick_start_triggered %d kick_start %d cur_matrix_base_addr %d vector_id %d",
                        //ram_addr, ram_cs, ram_we, ram_oe, kick_start_triggered, kick_start, cur_matrix_base_addr, vector_id);
                end
                READ_END: begin
                    //$display("kimi=====READ_END");
                    user_data_out <= ram_data;
                    ram_addr      <= 4'h0;
                    ram_cs        <= 1'b0;
                    ram_we        <= 1'b0;
                    ram_oe        <= 1'b0;
                    if (kick_start_triggered) begin
                        if (vector_id == 4'h3) begin
                            if (matrix_access_flag == MATRIX_A_READ) begin
                                state_machine  <= MAC_CACHE_LOAD_MAT_A_END;
                            end else if (matrix_access_flag == MATRIX_B_READ) begin
                                state_machine  <= MAC_CACHE_LOAD_MAT_B_END;
                            end else begin
                                state_machine  <= MAC_READ_AB_END;
                            end
                            //$display("[MATRIX READ COMPLETE]  matrix_access_flag %d vector_id: %d", matrix_access_flag, vector_id);
                        end else begin
                            state_machine <= READ_PRE;
                            vector_id     <= vector_id + 1;
                        end
                        if (cur_matrix_base_addr == MATRIX_A_BASE_ADDR) begin
                            matrix_a[vector_id] <= ram_data;
                            //$display("[READ_END] write matrix_a: base addr: 0x%h vector_id %d, data 0x%h", cur_matrix_base_addr, vector_id, ram_data);
                        end else if (cur_matrix_base_addr == MATRIX_B_BASE_ADDR) begin
                            matrix_b[vector_id] <= ram_data;
                            //$display("[READ_END] write matrix_b  base addr: 0x%h vector_id %d, data 0x%h", cur_matrix_base_addr, vector_id, ram_data);
                        end else begin
                            //$display("READ_END cur_matrix_base_addr %d", cur_matrix_base_addr);
                        end
                    end else begin
                        state_machine       <= MOD_BOOT;
                        user_data_out       <= ram_data;
                        ready               <= 1'b1;
                        kick_read_triggered <= 1'b0;
                    end
                    //$display("\n*************read READ_END addr: 0x%h ram_data: 0x%h cs: %d we: %d oe: %d", ram_addr, ram_data, ram_cs, ram_we, ram_oe);
                end
                /****************************** Mac Operation ************************************/
                MAC_START: begin
                    //$display("[MAC_START] kick_start_triggered: %d", kick_start_triggered);
                    state_machine  <= MAC_CACHE_LOAD_MAT_A_START;
                end
                MAC_CACHE_LOAD_MAT_A_START: begin
                    //$display("[MAC_CACHE_LOAD_MAT_A_START]=====");
                    cur_matrix_base_addr <= MATRIX_A_BASE_ADDR;
                    vector_id            <= 4'h0;
                    state_machine        <= READ_PRE;
                    matrix_access_flag   <= MATRIX_A_READ;
                end
                MAC_CACHE_LOAD_MAT_B_START: begin
                    //$display("[MAC_CACHE_LOAD_MAT_B_START]=====");
                    cur_matrix_base_addr <= MATRIX_B_BASE_ADDR;
                    vector_id            <= 4'h0;
                    state_machine        <= READ_PRE;
                    matrix_access_flag   <= MATRIX_B_READ;
                end
                MAC_READ_AB_END: begin
                    //$display("====[MAC_READ_AB_END]=====");
                    // 打印 matrix_a 的所有数据
                    $display("============================= Matrix A =============================");
                    for (integer i = 0; i < 4; i = i + 1) begin
                        $display("matrix_a[%0d] = 0x%h", i, matrix_a[i]);
                    end

                    // 打印 matrix_b 的所有数据  
                    $display("============================= Matrix B =============================");
                    for (integer i = 0; i < 4; i = i + 1) begin
                        $display("matrix_b[%0d] = 0x%h", i, matrix_b[i]);
                    end
                    vector_id          <= 4'h0;
                    state_machine      <= MAC_CALC_MAT_C;
                    //ready            <= 1'b1;
                    matrix_access_flag <= MATRIX_NONE;
                    user_data_out      <=32'h0;
                end
                MAC_CACHE_LOAD_MAT_A_END: begin
                    //$display("[MAC_CACHE_LOAD_MAT_A_END->MAC_CACHE_LOAD_MAT_B_START]================================");
                    state_machine      <= MAC_CACHE_LOAD_MAT_B_START;
                end
                MAC_CACHE_LOAD_MAT_B_END: begin
                    //$display("[MAC_CACHE_LOAD_MAT_B_END->MAC_READ_AB_END]=========================");
                    state_machine      <= MAC_READ_AB_END;
                end
                MAC_CALC_MAT_C: begin
                    $display("======================== MAC CALC Matrix C =========================");
                    matrix_mac(matrix_a, matrix_b, matrix_c);
                    state_machine      <= MAC_OUTPUT_MAT_C;
                end
                MAC_OUTPUT_MAT_C: begin
                    $display("============================= Matrix C =============================");
                    for (integer i = 0; i < 4; i = i + 1) begin
                        $display("matrix_c[%0d] = 0x%h", i, matrix_c[i]);
                    end
                    state_machine        <= MOD_POOLING_START;
                    ready                <= 1'b1;
                    kick_start_triggered <= 1'b0;
                    for (integer i = 0; i < 4; i = i + 1) begin
                        matrix_a_output[i] <= matrix_a[i];
                        matrix_b_output[i] <= matrix_b[i];
                        matrix_c_output[i] <= matrix_c[i];
                    end
                end
                MOD_POOLING_START: begin
                    $display("====================== POOLING CALC Matrix P =======================");
                    matrix_pooling(matrix_c, matrix_p);
                    $display("============================= Matrix P =============================");
                    $display("matrix_p[0] = 0x%h", matrix_p[31:16]);
                    $display("matrix_p[1] = 0x%h", matrix_p[15: 0]);
                    $display("matrix_p    = 0x%h", matrix_p);
                    state_machine <= MOD_POOLING_END;
                end
                MOD_POOLING_END: begin
                    state_machine   <= MOD_FUCKING_SLEEP;
                    matrix_p_output <= matrix_p;
                end
                default: begin
                    //$display("=====default=====");
                    vector_id          <= 4'h0;
                    state_machine      <= MOD_FUCKING_SLEEP;
                    matrix_access_flag <= MATRIX_NONE;
                end
            endcase
        end
    end

    // 实例化原始RAM
    single_port_sync_ram #(
        .ADDR_WIDTH(ADDR_WIDTH),
        .DATA_WIDTH(DATA_WIDTH),
        .DEPTH(DEPTH)
        ) ram_inst (
            .clk(clk),
            .addr(ram_addr),
            .data(ram_data),
            .cs(ram_cs),
            .we(ram_we),
            .oe(ram_oe) 
        );
endmodule